// MIT License
// Copyright (c) 2022 - 傅莘莘
// Source URL: https://github.com/zjhellofss/KuiperInfer
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

// Created by fss on 23-1-4.

#include <glog/logging.h>
#include <iostream>
#include <opencv2/opencv.hpp>
#include "../image_util.hpp"
#include "data/tensor.hpp"
#include "runtime/runtime_ir.hpp"
#include "tick.hpp"

//这个 PreProcessImage 函数的主要目的是将输入的图像转换成适合神经网络模型输入的格式。
//它首先调整图像大小，保持其纵横比并填充边缘以适应网络输入尺寸。
//然后，将图像转换为 RGB 格式并进行标准化处理。
//最后，它将处理后的图像数据转换为一个 Tensor 对象，以便可以被神经网络模型使用。
kuiper_infer::sftensor PreProcessImage(const cv::Mat& image, const int32_t input_h,
                                       const int32_t input_w) {
  //assert(!image.empty()); - 检查传入的图像是否为空。
  assert(!image.empty());
  using namespace kuiper_infer;
  //定义输入图像的通道数为 3，表示使用 RGB 颜色空间。
  const int32_t input_c = 3;
  //定义步长为 32，这可能用于图像的处理或缩放。
  int stride = 32;
  //创建一个用于存储处理后图像的 OpenCV 矩阵（Mat）对象。
  cv::Mat out_image;
  //调用 Letterbox 函数调整图像大小，使其适应给定的高度和宽度（input_h 和 input_w），同时保持图像的纵横比。填充的颜色为灰色（114, 114, 114）。
  Letterbox(image, out_image, {input_h, input_w}, stride, {114, 114, 114}, true);
  //创建一个新的 OpenCV 矩阵，用于存储转换为 RGB 颜色空间的图像。
  cv::Mat rgb_image;
  cv::cvtColor(out_image, rgb_image, cv::COLOR_BGR2RGB);

  cv::Mat normalize_image;
  //将 RGB 图像的像素值标准化到 0 到 1 的范围内。
  rgb_image.convertTo(normalize_image, CV_32FC3, 1. / 255.);

  std::vector<cv::Mat> split_images;
  cv::split(normalize_image, split_images);
  //检查分离后的图像通道数量是否正确。
  assert(split_images.size() == input_c);
  //创建一个新的 Tensor 对象，用于存储转换后的图像数据。
  std::shared_ptr<Tensor<float>> input = std::make_shared<Tensor<float>>(input_c, input_h, input_w);
  input->Fill(0.f);

  int index = 0;
  int offset = 0;
  //遍历分离的图像通道。
  for (const auto& split_image : split_images) {
    //确保每个通道的像素总数符合预期。
    assert(split_image.total() == input_w * input_h);
    const cv::Mat& split_image_t = split_image.t();
    //将通道数据复制到 Tensor 中。
    memcpy(input->slice(index).memptr(), split_image_t.data, sizeof(float) * split_image.total());
    index += 1;
    offset += split_image.total();
  }
  //返回处理后的 Tensor，它包含了预处理过的图像数据。
  return input;
}


// YoloDemo 函数展示了如何使用 YOLO 模型进行物体检测。它处理一批图像，对每张图像进行预处理，然后通过模型进行前向传播以获取检测结果。最后，它
// 将检测到的物体以边界框的形式绘制在图像上，并保存这些图像。此函数充分展示了从预处理图像到运行模型并获取结果，再到可视化这些结果的完整流程。
// 重要的是，这个函数也展示了如何使用非极大值抑制（NMS）技术来过滤重叠的边界框，确保每个检测到的物体只有一个边界框与之对应。这个过程中，每个检测都被分配了一个类别 ID 和置信度，这有助于理解模型如何识别和分类图像中的不同物体。
// 总的来说，YoloDemo 函数是一个展示如何在实际应用中使用 YOLO 模型进行物体检测的典型例子，它涵盖了从图像读取、预处理、模型推断到最终结果展示的整个流程。
void YoloDemo(const std::vector<std::string>& image_paths, const std::string& param_path,
              const std::string& bin_path, const uint32_t batch_size,
              const float conf_thresh = 0.25f, const float iou_thresh = 0.25f) {
  using namespace kuiper_infer;
  const int32_t input_h = 640;
  const int32_t input_w = 640;
  //创建 RuntimeGraph 对象，加载模型参数和二进制文件。
  RuntimeGraph graph(param_path, bin_path);
  //构建运行时图。
  graph.Build();

  assert(batch_size == image_paths.size());
  //创建一个存储输入张量的向量。
  std::vector<sftensor> inputs;
  // 遍历图像路径列表。
  for (uint32_t i = 0; i < batch_size; ++i) {
    //读取图像
    const auto& input_image = cv::imread(image_paths.at(i));
    sftensor input = PreProcessImage(input_image, input_h, input_w);
    //图像进行预处理。
    assert(input->rows() == 640);
    assert(input->cols() == 640);
    inputs.push_back(input);
  }
  //创建一个存储输出张量的向量。
  std::vector<std::shared_ptr<Tensor<float>>> outputs;
  //设置模型的输入。
  graph.set_inputs("pnnx_input_0", inputs);
  //执行一次前向传播。
  for (int i = 0; i < 1; ++i) {
    //运行模型进行前向传播。
    graph.Forward(true);
  }
  //获取模型的输出。
  outputs = graph.get_outputs("pnnx_output_0");
  assert(outputs.size() == inputs.size());
  assert(outputs.size() == batch_size);
  //遍历输出张量。
  for (int i = 0; i < outputs.size(); ++i) {
    //重新读取图像。
    const auto& image = cv::imread(image_paths.at(i));
    //获取原始图像的高度。
    const int32_t origin_input_h = image.size().height;
    const int32_t origin_input_w = image.size().width;
    //获取当前输出张量。
    const auto& output = outputs.at(i);
    assert(!output->empty());
    const auto& shapes = output->shapes();
    assert(shapes.size() == 3);

    const uint32_t elements = shapes.at(1);
    const uint32_t num_info = shapes.at(2);
    //创建一个存储检测结果的向量。
    std::vector<Detection> detections;

    std::vector<cv::Rect> boxes;
    std::vector<float> confs;
    std::vector<int> class_ids;

    const uint32_t b = 0;
    //遍历每个元素。
    for (uint32_t e = 0; e < elements; ++e) {
      // 对于每个元素，获取类别置信度、中心点坐标、宽度、高度等信息。
      // 根据这些信息计算边界框位置。
      // 确定每个边界框的最佳类别。
      float cls_conf = output->at(b, e, 4);
      if (cls_conf >= conf_thresh) {
        int center_x = (int)(output->at(b, e, 0));
        int center_y = (int)(output->at(b, e, 1));
        int width = (int)(output->at(b, e, 2));
        int height = (int)(output->at(b, e, 3));
        int left = center_x - width / 2;
        int top = center_y - height / 2;

        int best_class_id = -1;
        float best_conf = -1.f;
        for (uint32_t j = 5; j < num_info; ++j) {
          if (output->at(b, e, j) > best_conf) {
            best_conf = output->at(b, e, j);
            best_class_id = int(j - 5);
          }
        }

        boxes.emplace_back(left, top, width, height);
        confs.emplace_back(best_conf * cls_conf);
        class_ids.emplace_back(best_class_id);
      }
    }

    std::vector<int> indices;
    //应用非极大值抑制。
    cv::dnn::NMSBoxes(boxes, confs, conf_thresh, iou_thresh, indices);

    for (int idx : indices) {
      //对于每个选中的边界框，调整其坐标以适应原始图像大小。
      // 存储检测结果。
      Detection det;
      det.box = cv::Rect(boxes[idx]);
      ScaleCoords(cv::Size{input_w, input_h}, det.box, cv::Size{origin_input_w, origin_input_h});

      det.conf = confs[idx];
      det.class_id = class_ids[idx];
      detections.emplace_back(det);
    }

    int font_face = cv::FONT_HERSHEY_COMPLEX;
    double font_scale = 2;

    for (const auto& detection : detections) {
      //遍历所有检测到的物体。
      // 在原始图像上绘制边界框和类别ID。
      cv::rectangle(image, detection.box, cv::Scalar(255, 255, 255), 4);
      cv::putText(image, std::to_string(detection.class_id),
                  cv::Point(detection.box.x, detection.box.y), font_face, font_scale,
                  cv::Scalar(255, 255, 0), 4);
    }
    cv::imwrite(std::string("output") + std::to_string(i) + ".jpg", image);
  }
}

int main() {
  const uint32_t batch_size = 8;
  std::vector<std::string> image_paths;

  for (uint32_t i = 0; i < batch_size; ++i) {
    const std::string& image_path = "/code/KuiperInfer/imgs/bus.jpg";  // 可以放不同的图片
    image_paths.push_back(image_path);
  }
  const std::string& param_path = "/code/KuiperInfer/tmp/yolo/demo/yolov5s_batch8.pnnx.param";
  const std::string& bin_path = "/code/KuiperInfer/tmp/yolo/demo/yolov5s_batch8.pnnx.bin";

  YoloDemo(image_paths, param_path, bin_path, batch_size);
  return 0;
}
